#pragma once
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include <climits>
#include <iostream>
#include <ostream>
#include "HashPair.h"
#include "BAB_Lagrangian.h"
#include "assistant.h"
#include "k_median.h"

/*****************************************
 * 
 * 实现 BAB(branch-and-bound) 类
 * 类里面会存放一些全局信息，之后调用 run 方法，将运行分支定界算法
 * 
 *****************************************/
struct BAB
{
	std::unordered_map<int, std::unordered_map<int, double>>* distance;  // 记录距离矩阵
	std::vector<int>* current_id;  // 当前选出的设施 id
	std::vector<int>* fac_id;  // 所有设施的 id
	std::unordered_map<int, bool>* branch_flag;  // 记录是否在某个设施上进行了分支
	std::unordered_map<int, int>* visit_num;  // 记录每个设施被访问的数量
	std::unordered_map<int, char>* ftype;  // 设施 id 与设施类型之间的对应关系
	std::vector<HashPair<int>>* num2id;  // 记录数量到设施 id 的对应关系
	wh::reference_info_list* user_info;  // 记录用户的信息主要是为了获取 prob
	int k;  // 可以替换的设施数量
	int fnum;  // 原来的设施总数量
	double alpha;  // 原问题的上界
	double rate;  // 近似率
	double iter_num;  // 迭代次数
	std::ostream* fout; // 输出流的指针

	BAB()
	{
		distance = nullptr;
		current_id = nullptr;
		branch_flag = nullptr;
		visit_num = nullptr;
		ftype = nullptr;
		fac_id = nullptr;
		num2id = nullptr;
		fout = &std::cout;  // 默认初始化为 cout 
		k = INT_MIN;
		fnum = INT_MIN;
		alpha = DBL_MIN;
		rate = 0.02;
		iter_num = 10000;
	}

	/* 检查这个类是否进行初始化 */
	bool init()
	{
		return fac_id != nullptr && ftype != nullptr && distance != nullptr && current_id != nullptr && branch_flag != nullptr && visit_num != nullptr && num2id != nullptr && k != INT_MIN && fnum != INT_MIN && alpha != DBL_MIN;
	}

	/* 运行分支定界算法 */
	void run()
	{
		*fout << "问题的上界为：" << alpha << std::endl;
		std::unordered_set<int> selected;
		std::vector<int> unselected;
		for (auto iter = num2id->rbegin();iter != num2id->rend();iter++)
		{
			if ((*branch_flag)[iter->id] == false)
			{
				int f_id = iter->id;
				(*branch_flag)[f_id] = true;

				// 选中该设施
				selected.insert(f_id);
				run_BAB(selected, unselected);

				// 不选中该设施
				selected.erase(f_id);
				unselected.push_back(f_id);
				run_BAB(selected, unselected);
				unselected.pop_back();

				break;
			}
		}
	}
private:
	/* 递归地进行分支定界的算法，终止条件为所有分支都被剪掉 */
	void run_BAB(std::unordered_set<int>& selected, std::vector<int>& unselected)
	{
		/* 如果已经确定必须选择这些设施了，则不需要进行剪枝，可以直接确定精确值 */
		if (selected.size() == fnum)
		{
			double temp_dist = calculate_original_dist(*distance, selected);
			if(temp_dist < alpha)
			{
				current_id->clear();
				for (auto i : selected)
					current_id->push_back(i);
				alpha = calculate_original_dist(*distance, *current_id, *visit_num);  // 更新 alpha 和 visit_num
				num2id->clear();  // 更新 num2id
				for (auto p : (*visit_num))
					num2id->push_back(HashPair<int>(p));
				sort(num2id->begin(), num2id->end());
				*fout << "问题的新上界为：" << alpha << std::endl;
			}
			return;
		}

		/* 如果所有设施减去必须去除的设施，剩余数量等于需要的数量，也可以直接确定解 */
		if (fac_id->size() - unselected.size() == fnum)
		{
			std::unordered_set<int> unselected_set(unselected.begin(), unselected.end());
			std::vector<int> temp;
			for (auto i : *fac_id)
				if (unselected_set.count(i) == 0)
					temp.push_back(i);
			
			double temp_dist = calculate_original_dist(*distance, temp);
			if (temp_dist < alpha)
			{
				current_id->swap(temp);
				alpha = calculate_original_dist(*distance, *current_id, *visit_num);
				num2id->clear();  // 更新 num2id
				for (auto p : (*visit_num))
					num2id->push_back(HashPair<int>(p));
				sort(num2id->begin(), num2id->end());
				*fout << "问题的新上界为：" << alpha << std::endl;
			}
			return;
		}

		/* 正常情况下利用拉格朗日松弛进行剪枝 */
		auto r = Lagrangian(*distance, k, fnum, *ftype, *fac_id, alpha, selected, unselected, iter_num, rate, *fout, *user_info);
		
		/* 找到了更紧的上界，更新 */
		if (r.second != nullptr)
		{
			(*current_id).swap(*r.second);
			alpha = calculate_original_dist(*distance, *current_id, *visit_num);  // 更新上界
			num2id->clear();
			for (auto p : (*visit_num))
				num2id->push_back(HashPair<int>(p));
			sort(num2id->begin(), num2id->end());
			*fout << "问题的新上界为：" << alpha << std::endl;
			delete r.second;
		}

		/* 这个分支符合条件，进行下一步的分支 */
		if (r.first)
		{
			bool flag = true;
			/* 首先从已经选出的设施中确定可以进行分支的设施 */
			for (auto iter = num2id->rbegin(); iter != num2id->rend(); iter++)
			{
				if ((*branch_flag)[iter->id] == false)
				{
					int f_id = iter->id;
					(*branch_flag)[f_id] = true;
					flag = false;

					/* 首先测试选中它的情况 */
					selected.insert(f_id);
					run_BAB(selected, unselected);

					/* 测试未选中它的情况 */
					selected.erase(f_id);
					unselected.push_back(f_id);
					run_BAB(selected, unselected);
					unselected.pop_back();
					break;  // 只需找到一个分支即可
				}
			}

			/*****************************************
			 *
			 * 如果上述代码段未执行，则随机选择设施 id 进行分支
			 * 上述代码未成功执行的含义在于，对所有已经选出的点进行了分支，仍然不能剪掉所有分支
			 * 
			 *****************************************/
			if (flag)
			{
				for (int i : *fac_id)
				{
					if ((*branch_flag)[i] == false)
					{
						(*branch_flag)[i] = true;

						/* 首先测试选中它的情况 */
						selected.insert(i);
						run_BAB(selected, unselected);

						/* 测试未选中它的情况 */
						selected.erase(i);
						unselected.push_back(i);
						run_BAB(selected, unselected);
						unselected.pop_back();
						break;  // 只需找到一个分支即可
					}
				}
			}
		}
	}
};
